import pathlib
import typing

import torch

_M = typing.TypeVar("_M")
_R = typing.TypeVar("_R")
_O = typing.TypeVar("_O")


class Checkpoint(typing.Generic[_M, _R, _O]):
    def __init__(self, model_state: _M, random_state: _R, optimizer_state: _O):
        self._model_state = model_state
        self._random_state = random_state
        self._optimizer_state = optimizer_state

    def model_state(self) -> _M:
        return self._model_state

    def random_state(self) -> _R:
        return self._random_state

    def optimizer_state(self) -> _O:
        return self._optimizer_state

    def save(self, path: pathlib.Path) -> None:
        torch.save({
            "model": self._model_state,
            "random": self._random_state,
            "optimizer": self._optimizer_state
        }, path)

    @staticmethod
    def load(path: pathlib.Path) -> 'Checkpoint':
        checkpoint = torch.load(path)
        return Checkpoint(
            checkpoint["model"],
            checkpoint["random"],
            checkpoint["optimizer"]
        )
